"""
DTO生成器脚本
用于从YAML定义文件自动生成C++ DTO（数据传输对象）头文件
支持Qt类型、可选字段、嵌套DTO结构等特性
"""

import yaml
import os
import argparse
from datetime import datetime
import collections

# Qt类型及其对应的头文件包含
QT_TYPES_INCLUDES = {
    "QString": "<QString>",
    "QStringList": "<QStringList>",
    "QDateTime": "<QDateTime>",
    "QDate": "<QDate>",
    "QTime": "<QTime>",
    "QVariant": "<QVariant>",
    "QByteArray": "<QByteArray>",
    "QList": "<QList>",
    "QVector": "<QVector>",
    "QMap": "<QMap>",
    "QUuid": "<QUuid>",
}

# 基本数据类型列表
PRIMITIVE_TYPES = ["bool", "int", "float", "double", "char", "void",
                   "qint8", "qint16", "qint32", "qint64",
                   "quint8", "quint16", "quint32", "quint64", "size_t"]

# 全局存储：所有DTO定义（名称 -> 定义）和DTO名称到输出文件详情的映射
ALL_DTO_DEFINITIONS = {}
DTO_TO_FILE_MAP = {}


def get_cpp_type_info(field_type, is_optional):
    """
    确定C++类型信息，包括是否为指针类型（用于可选DTO）
    返回元组：(cpp_type_str, is_pointer_type, base_type_if_dto_or_none)
    
    参数:
        field_type: 字段类型字符串
        is_optional: 是否为可选字段
    
    返回:
        (cpp类型字符串, 是否为指针类型, 如果是DTO则返回基础类型否则为None)
    """
    base_type = field_type
    is_pointer_type = False
    base_type_if_dto = None

    # 处理容器类型：QList<T>, QVector<T>, QMap<K,V>
    if field_type.startswith("QList<") and field_type.endswith(">"):
        inner_type_name = field_type[6:-1].strip()
        base_type = f"QList<{inner_type_name}>"
        if inner_type_name in ALL_DTO_DEFINITIONS:
             pass # QList<MyDTO> - MyDTO将由其自己的定义处理
    elif field_type.startswith("QVector<") and field_type.endswith(">"):
        base_type = field_type
    elif field_type.startswith("QMap<") and field_type.endswith(">"):
        base_type = field_type
    elif base_type in ALL_DTO_DEFINITIONS: # 引用其他DTO
        base_type_if_dto = base_type
        if is_optional:
            # 对可选DTO使用QSharedPointer（有助于前向声明和堆管理）
            base_type = f"QSharedPointer<{base_type}>"
            is_pointer_type = True
        # else: 非可选DTO成员按值传递（base_type保持为DTO名称）

    # 对非DTO可选类型使用std::optional
    if is_optional and not is_pointer_type and base_type_if_dto is None:
        base_type = f"std::optional<{base_type}>"

    return base_type, is_pointer_type, base_type_if_dto


def generate_cpp_for_file_group(file_group_def, all_dto_definitions_map, dto_to_file_map):
    """
    为单个输出文件组生成C++代码
    
    参数:
        file_group_def: 文件组定义（包含输出路径、文件名、DTO列表等）
        all_dto_definitions_map: 所有DTO定义的映射
        dto_to_file_map: DTO名称到文件详情的映射
    """
    output_path_dir = file_group_def['output_path']
    output_filename = file_group_def['filename']
    dtos_in_group = file_group_def['dtos']
    current_namespace = file_group_def.get('namespace')

    # 确保输出目录存在
    if not os.path.exists(output_path_dir):
        try:
            os.makedirs(output_path_dir)
            print(f"创建输出目录: {output_path_dir}")
        except OSError as e:
            print(f"错误: 无法创建输出目录 '{output_path_dir}': {e}")
            return None
    
    output_file_full_path = os.path.join(output_path_dir, output_filename)

    cpp_code = []
    generation_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC")

    # 添加文件头注释
    cpp_code.append(f"// 由DTO生成器脚本生成于 {generation_timestamp}")
    cpp_code.append(f"// 目标文件: {output_filename}")
    cpp_code.append("#pragma once")
    cpp_code.append("")

    # 收集必要的头文件包含
    local_includes = set()
    needs_optional_for_this_file = False
    needs_qsharedpointer_for_this_file = False
    cross_file_dto_includes = set() # 存储需要包含的其他生成的.h文件

    # 本文件中定义的DTO的前向声明
    dto_names_in_this_file = {dto_def['name'] for dto_def in dtos_in_group}

    # 处理每个DTO的字段，收集必要的头文件包含
    for dto_def in dtos_in_group:
        if 'fields' in dto_def:
            for field in dto_def['fields']:
                yaml_field_type = field['type']
                is_optional = field.get('optional', False)

                cpp_type_str, is_pointer, base_dto_type = get_cpp_type_info(yaml_field_type, is_optional)

                # 添加Qt类型的头文件包含
                if yaml_field_type in QT_TYPES_INCLUDES:
                    local_includes.add(QT_TYPES_INCLUDES[yaml_field_type])
                elif yaml_field_type.startswith("QList<"):
                    local_includes.add("<QList>")
                    inner = yaml_field_type[6:-1].strip()
                    if inner in QT_TYPES_INCLUDES: local_includes.add(QT_TYPES_INCLUDES[inner])
                elif yaml_field_type.startswith("QVector<"):
                    local_includes.add("<QVector>")
                    inner = yaml_field_type[8:-1].strip()
                    if inner in QT_TYPES_INCLUDES: local_includes.add(QT_TYPES_INCLUDES[inner])
                elif yaml_field_type.startswith("QMap<"):
                    local_includes.add("<QMap>")
                    parts = yaml_field_type[5:-1].split(',', 1)
                    if parts[0].strip() in QT_TYPES_INCLUDES: local_includes.add(QT_TYPES_INCLUDES[parts[0].strip()])
                    if len(parts) > 1 and parts[1].strip() in QT_TYPES_INCLUDES: local_includes.add(QT_TYPES_INCLUDES[parts[1].strip()])
                
                # 检查是否需要std::optional或QSharedPointer
                if "std::optional<" in cpp_type_str:
                    needs_optional_for_this_file = True
                if "QSharedPointer<" in cpp_type_str:
                    needs_qsharedpointer_for_this_file = True

                # 检查字段类型是否是在其他文件中定义的DTO
                if base_dto_type and base_dto_type not in dto_names_in_this_file:
                    if base_dto_type in dto_to_file_map:
                        other_dto_file_details = dto_to_file_map[base_dto_type]
                        if other_dto_file_details["filename"] != output_filename:
                             cross_file_dto_includes.add(f"\"{other_dto_file_details['filename']}\"")
                    else:
                        print(f"警告: DTO类型 '{base_dto_type}' 在 '{dto_def['name']}' (文件: {output_filename}) 中使用，但在dto_to_file_map中未找到其定义文件。")

    # 添加必要的标准库头文件
    if needs_optional_for_this_file:
        local_includes.add("<optional>")
    if needs_qsharedpointer_for_this_file:
        local_includes.add("<QSharedPointer>")
    
    # 添加所有头文件包含
    sorted_local_includes = sorted(list(local_includes))
    for inc in sorted_local_includes:
        cpp_code.append(f"#include {inc}")
    
    # 添加其他DTO文件的包含
    sorted_cross_file_includes = sorted(list(cross_file_dto_includes))
    if sorted_cross_file_includes:
        cpp_code.append("")
        cpp_code.append("// 包含在其他文件中定义的DTO")
        for inc in sorted_cross_file_includes:
            cpp_code.append(f"#include {inc}")

    cpp_code.append("")
    cpp_code.append("#include <QMetaType> // 用于Q_DECLARE_METATYPE")
    cpp_code.append("")

    # 添加命名空间（如果指定）
    if current_namespace:
        cpp_code.append(f"namespace {current_namespace} {{")
        cpp_code.append("")

    # 添加本文件中DTO的前向声明
    cpp_code.append(f"// 本文件中DTO的前向声明 ({output_filename})")
    for dto_def in dtos_in_group:
        cpp_code.append(f"struct {dto_def['name']};")
    cpp_code.append("")

    # 生成每个DTO的结构体定义
    for dto_def in dtos_in_group:
        dto_name = dto_def['name']
        if dto_def.get('description'):
            cpp_code.append(f"/**\n * @brief {dto_def['description']}\n */")
        cpp_code.append(f"struct {dto_name} {{")

        constructor_params = []
        constructor_body = []
        member_declarations = []

        if 'fields' in dto_def and dto_def['fields']:
            for field in dto_def['fields']:
                field_name = field['name']
                yaml_field_type = field['type']
                is_optional = field.get('optional', False)
                
                cpp_type, is_pointer_type, _ = get_cpp_type_info(yaml_field_type, is_optional)
                
                field_desc = field.get('description', '')
                member_declarations.append(f"    {cpp_type} {field_name};{(' // ' + field_desc) if field_desc else ''}")
                
                param_type = cpp_type
                default_value = ""
                is_dto_ref_type = yaml_field_type in ALL_DTO_DEFINITIONS

                # 设置构造函数参数的默认值
                if is_optional:
                    if is_pointer_type: # QSharedPointer<DTO>
                        default_value = " = {}" 
                    else: # std::optional<NonDTOType>
                        param_type = f"const {cpp_type}&"
                        default_value = " = std::nullopt"
                elif is_dto_ref_type: # 非可选DTO成员
                     param_type = f"const {cpp_type}&"
                     default_value = " = {}"
                else: # 非可选基本类型或Qt值类型
                    if cpp_type == "bool": default_value = " = false"
                    elif cpp_type in PRIMITIVE_TYPES and "q" in cpp_type or "int" in cpp_type or "size_t" == cpp_type : default_value = " = 0"
                    elif cpp_type in ["float", "double"]: default_value = " = 0.0"
                    elif cpp_type in ["QString", "QStringList", "QByteArray", "QDate", "QTime", "QUuid"] or \
                         cpp_type.startswith("QList<") or cpp_type.startswith("QVector<") or \
                         cpp_type.startswith("QMap<") or cpp_type == "QDateTime" or cpp_type == "QVariant":
                        default_value = " = {}" 
                
                constructor_params.append(f"{param_type} _{field_name}{default_value}")
                constructor_body.append(f"        {field_name}(_{field_name})")
            
            # 添加成员变量声明
            for decl in member_declarations:
                 cpp_code.append(decl)
            cpp_code.append("")

            # 添加默认构造函数
            cpp_code.append(f"    {dto_name}() = default;")
            cpp_code.append("")
            
            # 添加带参数的构造函数
            if constructor_params:
                cpp_code.append(f"    explicit {dto_name}(")
                for i, param_str in enumerate(constructor_params):
                    cpp_code.append(f"        {param_str}{',' if i < len(constructor_params) - 1 else ''}")
                cpp_code.append("    )")
                if constructor_body: cpp_code.append("    : " + "\n    , ".join(constructor_body))
                cpp_code.append("    {}")
                cpp_code.append("")
        else:
            cpp_code.append(f"    // {dto_name}没有定义字段");
            cpp_code.append(f"    {dto_name}() = default;");
        cpp_code.append("};")
        cpp_code.append("")

    # 关闭命名空间
    if current_namespace:
        cpp_code.append(f"}} // namespace {current_namespace}")
        cpp_code.append("")

    # 添加QMetaType声明
    cpp_code.append(f"// {output_filename}中DTO的元类型声明")
    for dto_def in dtos_in_group:
        dto_name = dto_def['name']
        fq_dto_name = f"{current_namespace}::{dto_name}" if current_namespace else dto_name
        cpp_code.append(f"Q_DECLARE_METATYPE({fq_dto_name})")
    
    cpp_code.append("")
    cpp_code.append(f"// {output_filename}的DTO生成结束")

    # 写入文件
    try:
        with open(output_file_full_path, 'w', encoding='utf-8') as f:
            f.write("\n".join(cpp_code))
        print(f"成功生成DTO头文件: {output_file_full_path}")
    except IOError as e:
        print(f"写入输出文件时出错 '{output_file_full_path}': {e}")
    except Exception as e:
        print(f"写入 '{output_file_full_path}' 时发生意外错误: {e}")


def main_generate_all(yaml_file_path):
    """
    主函数：从YAML文件生成所有DTO
    
    参数:
        yaml_file_path: YAML定义文件的路径
    """
    try:
        with open(yaml_file_path, 'r', encoding='utf-8') as f:
            data = yaml.safe_load(f)
    except FileNotFoundError:
        print(f"错误: 找不到YAML文件 '{yaml_file_path}'.")
        return
    except yaml.YAMLError as e:
        print(f"解析YAML文件时出错 '{yaml_file_path}': {e}")
        return
    except Exception as e:
        print(f"读取 '{yaml_file_path}' 时发生意外错误: {e}")
        return

    if not data or 'dto_files' not in data:
        print("错误: YAML文件必须包含'dto_files'列表，其中每个项目定义输出文件及其DTO。")
        return

    # 第一遍：收集所有DTO定义并映射DTO名称到其目标文件详情
    global ALL_DTO_DEFINITIONS, DTO_TO_FILE_MAP
    for file_group in data['dto_files']:
        if 'dtos' in file_group and isinstance(file_group['dtos'], list):
            for dto_def in file_group['dtos']:
                dto_name = dto_def['name']
                if dto_name in ALL_DTO_DEFINITIONS:
                    print(f"警告: DTO名称 '{dto_name}' 被多次定义。将使用最后找到的定义。")
                ALL_DTO_DEFINITIONS[dto_name] = dto_def
                DTO_TO_FILE_MAP[dto_name] = {
                    "filename": file_group['filename'],
                    "output_path": file_group['output_path'],
                    "namespace": file_group.get('namespace')
                }
        else:
            print(f"警告: 文件组 '{file_group.get('filename', '未知文件')}' 没有'dtos'列表或不是列表。跳过。")

    # 第二遍：生成每个文件
    for file_group_def in data['dto_files']:
        if 'output_path' not in file_group_def or 'filename' not in file_group_def or 'dtos' not in file_group_def:
            print(f"警告: 由于缺少'output_path'、'filename'或'dtos'键，跳过文件组。详情: {file_group_def}")
            continue
        print(f"\n处理文件DTO: {os.path.join(file_group_def['output_path'], file_group_def['filename'])}")
        generate_cpp_for_file_group(file_group_def, ALL_DTO_DEFINITIONS, DTO_TO_FILE_MAP)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="从YAML定义文件生成C++ DTO头文件，该文件按输出文件对DTO进行分组。")
    parser.add_argument("yaml_file", help="输入YAML DTO定义文件的路径。")
    
    args = parser.parse_args()
    main_generate_all(args.yaml_file)